// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: Chunhui

#pragma once

#include <memory>
#include <string>
#include <stdexcept>
#include <vector>

#include <glog/logging.h>
#include <highfive/H5File.hpp>
#include <highfive/H5DataSet.hpp>
#include <highfive/H5Group.hpp>

namespace interval_feed {

class TimedDataUniverseFixedDumper {
 public:
  TimedDataUniverseFixedDumper(
      const std::string& filepath, uint64_t nuniv, uint64_t nrow)
      : nuniv_(nuniv),
        nrow_(nrow),
        ncol_(0) {
    file_ = std::make_unique<HighFive::File>(
          filepath,
          HighFive::File::ReadWrite | HighFive::File::Create | HighFive::File::Truncate);
  }

  uint64_t nuniv() const {
    return nuniv_;
  }
  uint64_t ncol() const {
    return ncol_;
  }
  uint64_t nrow() const {
    return nrow_;
  }

  void set_universe(const std::vector<std::string>& universe) {
    CHECK_EQ(nuniv_, universe.size());
    write_dataset<std::string>("universe", universe);
  }

  void set_session_names(const std::vector<std::string>& session_names) {
    CHECK_EQ(nrow_, session_names.size());
    write_dataset<std::string>("session_names", session_names);
  }

  void set_timestamp(const std::vector<int64_t>& timestamp) {
    CHECK_EQ(nrow_, timestamp.size());
    write_dataset<int64_t>("timestamp", timestamp);
  }

  template <typename T>
  void set_column(const std::string& column_name, const std::vector<T>& column) {
    CHECK_EQ(nrow_ * nuniv_, column.size());
    std::vector<size_t> dims = {nrow_, nuniv_};
    write_dataset<T>(column_name, column, dims);
    ncol_++;
  }

  void set_column_float(const std::string& column_name, const std::vector<float>& column) {
    CHECK_EQ(nrow_ * nuniv_, column.size());
    std::vector<size_t> dims = {nrow_, nuniv_};
    write_dataset<float>(column_name, column, dims);
    ncol_++;
  }

  void set_column_double(const std::string& column_name, const std::vector<double>& column) {
    CHECK_EQ(nrow_ * nuniv_, column.size());
    std::vector<size_t> dims = {nrow_, nuniv_};
    write_dataset<double>(column_name, column, dims);
    ncol_++;
  }

 protected:
  const uint64_t nuniv_;
  const uint64_t nrow_;
  std::unique_ptr<HighFive::File> file_;
  uint64_t ncol_;

  template <typename T>
  void write_dataset(
        const std::string& name, const std::vector<T>& data, const std::vector<size_t>& dims) {
    HighFive::DataSetCreateProps props;
    std::vector<hsize_t> chunk_dims(dims.begin(), dims.end());
    props.add(HighFive::Chunking(chunk_dims));
    props.add(HighFive::Deflate(7));
    CHECK_LE(dims.size(), 2);
    CHECK_GE(dims.size(), 1);

    HighFive::DataSet dataset = file_->createDataSet<T>(
        name,
        HighFive::DataSpace(dims),
        props,
        HighFive::DataSetAccessProps());

    if (dims.size() == 2) {
      // TODO(chunhui): inefficient. change to write_raw with coerced 2d array.
      std::vector<std::vector<T>> data_2d;
      reshape_2d_data(data, dims, &data_2d);
      dataset.write(data_2d);
    } else {
      dataset.write(data);
    }
  }

  template <typename T>
  void reshape_2d_data(const std::vector<T>& data, const std::vector<size_t>& dims,
      std::vector<std::vector<T>>* data_2d) {
    data_2d->clear();
    size_t nrows = dims[0];
    size_t ncols = dims[1];
    for (size_t i = 0; i < nrows; ++i) {
      const auto start = data.begin() + i * ncols;
      data_2d->emplace_back(start, start + ncols);
    }
  }

  template <typename T>
  void write_dataset(const std::string& name, const std::vector<T>& data) {
    std::vector<size_t> dims = {data.size()};
    write_dataset<T>(name, data, dims);
  }
};

}  // namespace interval_feed
